import os
import cv2
import numpy as np
from skimage import morphology

def i2f(i_image):
    f_image = np.float32(i_image) / 255.0
    return f_image

def f2i(f_image):
    i_image = np.uint8(f_image * 255.0)
    return i_image


#使用直方图统计和中值滤波来估计大气光，避免了因局部亮度过高导致的误判
def Compute_A_Tang(im, window_size=15):
    R = im[:, :, 2]
    G = im[:, :, 1]
    B = im[:, :, 0]

    # 自适应窗口大小的暗通道计算
    dark = np.min(im, axis=2)
    dark = cv2.medianBlur(dark.astype(np.uint8), window_size)

    # 基于暗通道先验估计大气光
    [h, edges] = np.histogram(dark, 200)
    numpixel = im.shape[0] * im.shape[1]
    thr_frac = numpixel * 0.99
    csum = np.cumsum(h)
    nz_idx = np.nonzero(csum > thr_frac)[0][0]
    dc_thr = edges[nz_idx]
    mask = dark >= dc_thr

    rs = R[mask]
    gs = G[mask]
    bs = B[mask]

    A = np.zeros((1, 3))
    A[0, 2] = np.median(rs)
    A[0, 1] = np.median(gs)
    A[0, 0] = np.median(bs)

    return A

def GetIntensity(fi):
    return cv2.divide(fi[:, :, 0] + fi[:, :, 1] + fi[:, :, 2], 3)

def GetSaturation(fi, intensity):
    min_rgb = cv2.min(cv2.min(fi[:, :, 0], fi[:, :, 1]), fi[:, :, 2])
    me = np.finfo(np.float32).eps
    S = 1.0 - min_rgb / (intensity + me)
    return S

def EstimateSaturation_Gamma(h_saturation, g=0.2):
    j_saturation = (np.power(h_saturation, 1.0 / g) + 1.0 - np.power(1.0 - h_saturation, 1.0 / g)) / 2.0
    j_saturation = np.maximum(j_saturation, h_saturation)
    return j_saturation

def EstimateTransimission(h_intensity, h_saturation, j_saturation):
    Td = h_intensity * (j_saturation - h_saturation)
    Tmn = j_saturation
    Tmap = 1.0 - (Td / Tmn)
    me = np.finfo(np.float32).eps
    Tmap = np.clip(Tmap, me, 1.0)

    # 使用引导滤波平滑透射率
    Tmap = cv2.ximgproc.guidedFilter(h_intensity.astype(np.float32), Tmap.astype(np.float32), 15, 0.2)
    return Tmap

def Recover(im, tmap, A, min_trans=0.1):
    res = np.empty(im.shape, im.dtype)
    for ind in range(0, 3):
        # 限制透射率的最小值
        res[:, :, ind] = (im[:, :, ind] - A[0, ind]) / np.maximum(tmap, min_trans) + A[0, ind]
        res[:, :, ind] = np.clip(res[:, :, ind], 0.0, 1.0)
    return res

def Adjust(im, perh=99.9, perl=0.5):
    aim = np.empty(im.shape, im.dtype)
    temp = np.empty(im.shape, im.dtype)
    im_h = np.percentile(im, perh)
    im_l = np.percentile(im, perl)
    for ind in range(0, 3):
        aim[:, :, ind] = (im[:, :, ind] - im_l) / (im_h - im_l)
        temp[:, :, ind] = np.clip(aim[:, :, ind], 0.0, 1.0)
    return aim

def Normalize(im):
    aim = np.empty(im.shape, im.dtype)
    for ind in range(0, 3):
        im_h = np.max(im[:, :, ind])
        im_l = np.min(im[:, :, ind])
        aim[:, :, ind] = (im[:, :, ind] - im_l) / (im_h - im_l)
        aim[:, :, ind] = np.clip(aim[:, :, ind], 0.0, 1.0)
    return aim

def gray_world(im):
    aim = np.empty(im.shape, im.dtype)
    mu_r = np.average(im[:, :, 2])
    mu_g = np.average(im[:, :, 1])
    mu_b = np.average(im[:, :, 0])
    aim[:, :, 0] = np.minimum(im[:, :, 0] * (mu_g / mu_b), 1.0)
    aim[:, :, 2] = np.minimum(im[:, :, 2] * (mu_g / mu_r), 1.0)
    aim[:, :, 1] = im[:, :, 1]
    return aim

def Clahe(im, clip=1):
    HSV = cv2.cvtColor(f2i(im), cv2.COLOR_BGR2HSV)
    clahe = cv2.createCLAHE(clipLimit=clip, tileGridSize=(8, 8))
    HSV[:, :, 2] = clahe.apply(HSV[:, :, 2])
    result_im = i2f(cv2.cvtColor(HSV, cv2.COLOR_HSV2BGR))
    return result_im

def process_image(input_path, output_dir):
    hazy_image = i2f(cv2.imread(input_path, cv2.IMREAD_COLOR))

    # 降低去雾效果
    A = Compute_A_Tang(hazy_image, window_size=15)
    A = A * 0.8  # 降低大气光估计值，以降低去雾效果

    # 白平衡处理
    hazy_imageWB = gray_world(hazy_image)
    A_WB = Compute_A_Tang(hazy_imageWB, window_size=15)
    A_WB = A_WB * 0.8  # 降低大气光估计值，以降低去雾效果

    S_A = np.max(A) - np.min(A)
    S_AWB = np.max(A_WB) - np.min(A_WB)

    perh = 99.9
    perl = 0.5
    epsilon = 0.02
    cl = 1

    if S_A < S_AWB + epsilon:
        hazy_imagen = np.empty(hazy_image.shape, hazy_image.dtype)
        for ind in range(0, 3):
            hazy_imagen[:, :, ind] = hazy_image[:, :, ind] / A[0, ind]
        hazy_imagen = Normalize(hazy_imagen)
        hazy_I = GetIntensity(hazy_imagen)
        hazy_S = GetSaturation(hazy_imagen, hazy_I)
        est_S = EstimateSaturation_Gamma(hazy_S, 0.2)
        Transmap = EstimateTransimission(hazy_I, hazy_S, est_S)
        r_image = Recover(hazy_image, Transmap, A, min_trans=0.1)
        r_image = Adjust(r_image, perh, perl)
    else:
        hazy_imagen = np.empty(hazy_image.shape, hazy_image.dtype)
        for ind in range(0, 3):
            hazy_imagen[:, :, ind] = hazy_image[:, :, ind] / A_WB[0, ind]
        hazy_imagen = Normalize(hazy_imagen)
        hazy_I = GetIntensity(hazy_imagen)
        hazy_S = GetSaturation(hazy_imagen, hazy_I)
        est_S = EstimateSaturation_Gamma(hazy_S, 0.2)
        Transmap = EstimateTransimission(hazy_I, hazy_S, est_S)
        r_image = Recover(hazy_image, Transmap, A_WB, min_trans=0.1)
        r_image = Adjust(r_image, perh, perl)
        r_image = gray_world(r_image)

    # 降低CLAHE的clipLimit
    result_ce = Clahe(r_image, cl)

    # 添加颜色校正步骤
    result_ce = cv2.cvtColor((result_ce * 255).astype(np.uint8), cv2.COLOR_BGR2LAB)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    result_ce[:, :, 0] = clahe.apply(result_ce[:, :, 0])
    result_ce = cv2.cvtColor(result_ce, cv2.COLOR_LAB2BGR)
    result_ce = i2f(result_ce)

    # 减少彩色噪点
    result_ce = cv2.medianBlur(result_ce, 3)
    result_ce = cv2.bilateralFilter(result_ce, 9, 75, 75)

    filename = os.path.basename(input_path)
    output_filename = f"processed_{filename}"
    output_path = os.path.join(output_dir, output_filename)
    cv2.imwrite(output_path, f2i(result_ce))

    return output_filename